import os
import sys

from flask import Flask, request, jsonify
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores.milvus import Milvus
from pymilvus import connections
from pymilvus.orm import utility
from sentence_transformers import SentenceTransformer

cur_path = os.path.abspath(os.path.dirname(__file__))
sys.path.insert(0, cur_path + "/..")

from src.config import MILVUS_HOST, MILVUS_PORT, SERVER_HOST, SERVER_PORT, MILVUS_DB, MILVUS_USER, MILVUS_PASSWORD
from src.entity.M3eEmbdddings import M3eEmbeddings

# 模型加载
model = SentenceTransformer(cur_path + '/m3e-base')
# 连接milvus
connections.connect(host=MILVUS_HOST, port=MILVUS_PORT, db_name=MILVUS_DB, alias='milvus_connection'
                    , user=MILVUS_USER, password=MILVUS_PASSWORD)
milvus_store = Milvus(
    embedding_function=M3eEmbeddings(model=model),
    vector_field='vector',
    primary_field='id',
    connection_args={"host": MILVUS_HOST, "port": MILVUS_PORT, "db_name": MILVUS_DB, "user": MILVUS_USER,
                     "password": MILVUS_PASSWORD}
)

app = Flask(__name__)
# app.config['JSON_AS_ASCII'] = False
app.json.ensure_ascii = False # 解决中文乱码问题
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=50,
    separators=["\n\n", "\n", "\. ", " ", ""]
)


# 创建集合并添加数据
@app.route('/create_knowledge_db', methods=['POST'])
def create_knowledge_db_controller():
    contents = request.json.get('contents')
    collection_name = request.json.get('collection_name')
    # # 创建集合并加载
    # create_knowledge_db(collection_name)
    # 加载数据
    text_chunks = []
    for content in contents:
        text_chunks.extend(text_splitter.split_text(content))
    mata_data = [{"content": content_chunk} for content_chunk in text_chunks]
    milvus_store.from_texts(text_chunks,
                            collection_name=collection_name,
                            metadatas=mata_data,
                            embedding=M3eEmbeddings(model=model),
                            connection_args={"host": MILVUS_HOST, "port": MILVUS_PORT, "db_name": MILVUS_DB,
                                             "user": MILVUS_USER, "password": MILVUS_PASSWORD}
                            )
    return jsonify({
        'code': 200,
        'message': '创建成功',
        'data': collection_name
    })


@app.route('/drop_knowledge_db', methods=['POST'])
def drop_knowledge_db_controller():
    collection_name = request.json.get('collection_name')
    utility.drop_collection(collection_name, using='milvus_connection')
    return jsonify({
        'code': 200,
        'message': '删除成功',
        'data': collection_name
    })


@app.route('/search-collection-list', methods=['GET'])
def get_collection_list_controller():
    collection_names = utility.list_collections(using='milvus_connection')
    return jsonify({
        'code': 200,
        'message': '查询成功',
        'data': collection_names
    })


@app.route('/search-content', methods=['POST'])
def search_content():
    keyword = request.json.get('keyword')
    search_limit = request.json.get('limit')
    collection_name = request.json.get('collection_name')
    milvus = Milvus(
        embedding_function=M3eEmbeddings(model=model),
        collection_name=collection_name,
        connection_args={"host": MILVUS_HOST, "port": MILVUS_PORT, "db_name": MILVUS_DB, "user": MILVUS_USER,
                         "password": MILVUS_PASSWORD}
    )
    results = milvus.similarity_search(keyword, k=search_limit)
    result_list = [{'content': result.page_content} for result in results]
    return jsonify({
        'code': 200,
        'message': '查询成功',
        'data': result_list
    })


if __name__ == '__main__':
    app.run(host=SERVER_HOST, port=SERVER_PORT)
